package cn.fetosoft.woodpecker.core.jmeter;

import cn.fetosoft.woodpecker.core.data.entity.BaseElement;
import cn.fetosoft.woodpecker.core.data.entity.element.TestPlanElement;
import cn.fetosoft.woodpecker.core.data.entity.element.sampler.HttpSamplerElement;
import cn.fetosoft.woodpecker.core.data.form.ElementForm;
import cn.fetosoft.woodpecker.core.data.service.ElementService;
import cn.fetosoft.woodpecker.core.enums.ElementCategory;
import cn.fetosoft.woodpecker.core.event.BuildTreeEvent;
import cn.fetosoft.woodpecker.core.jmeter.extension.HttpHeaderConfig;
import cn.fetosoft.woodpecker.core.jmeter.extension.ModuleControllerExt;
import org.apache.jmeter.JMeter;
import org.apache.jmeter.engine.StandardJMeterEngine;
import org.apache.jmeter.protocol.http.control.Header;
import org.apache.jmeter.protocol.http.control.HeaderManager;
import org.apache.jmeter.protocol.http.sampler.HttpRequestSampler;
import org.apache.jmeter.save.SaveService;
import org.apache.jmeter.testelement.TestElement;
import org.apache.jmeter.testelement.property.CollectionProperty;
import org.apache.jmeter.testelement.property.JMeterProperty;
import org.apache.jmeter.testelement.property.PropertyIterator;
import org.apache.jmeter.threads.AbstractThreadGroup;
import org.apache.jmeter.threads.PostThreadGroup;
import org.apache.jmeter.util.JMeterUtils;
import org.apache.jorphan.collections.ListedHashTree;
import org.apache.jorphan.collections.SearchByClass;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import java.io.File;
import java.io.FileOutputStream;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.stream.Collectors;

/**
 * @author guobingbing
 * @version 1.0
 * @wechat t_gbinb
 * @create 2021/6/21 15:24
 */
@Component
public class JMeterBuilder implements JMeterService, ApplicationContextAware {

	private static final Logger log = LoggerFactory.getLogger(JMeterBuilder.class);
	private static final String USER_DIR = System.getProperty("user.dir");
	private static final String JMETER_PATH =  USER_DIR + "/jmeter";
	private static final String LIB_EXT_PATH = USER_DIR + "/lib";

	@Autowired
	private ElementService elementService;
	private ApplicationContext applicationContext;

	static {
		System.out.println("=============================");
		System.out.println(JMETER_PATH);
		System.out.println("=============================");
		// 设置不适用gui的方式调用jmeter
		System.setProperty(JMeter.JMETER_NON_GUI, "true");
		File file = new File(JMETER_PATH + "/jmeter.properties");
		JMeterUtils.loadJMeterProperties(file.getAbsolutePath());
		JMeterUtils.setProperty("search_paths", LIB_EXT_PATH);
		JMeterUtils.setLocale(Locale.SIMPLIFIED_CHINESE);
		//加载扩展jar
		LoadExtLibsUtil.loadJars(LIB_EXT_PATH);
	}

	@Override
	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException{
		this.applicationContext = applicationContext;
	}

	@Override
	public void startTest(String userId, String testId, String planId) throws Exception {
		ElementForm form = new ElementForm();
		form.setCategory(ElementCategory.TEST_PLAN.getName());
		form.setPlanId(planId);
		List<TestPlanElement> planList = elementService.selectListByForm(form, TestPlanElement.class);
		if(!CollectionUtils.isEmpty(planList)){
			if(planList.size()>1){
				throw new Exception("TestPlan more than one!!!");
			}
			TestPlanElement testPlanElement = planList.get(0);
			testPlanElement.setUserId(userId);
			testPlanElement.setTestId(testId);
			// jemter 引擎
			StandardJMeterEngine engine = new StandardJMeterEngine();
			ListedHashTree rootTree = this.buildJMeterEngine(testPlanElement);
			engine.configure(rootTree);
			engine.runTest();
		}else{
			throw new Exception("TestPlan is empty!!!");
		}
	}

	private ListedHashTree buildJMeterEngine(TestPlanElement testPlanElement) throws Exception{
		ElementForm form = new ElementForm();
		form.setPlanId(testPlanElement.getPlanId());
		form.setRows(0);
		form.setAscField("createTime");
		List<BaseElement> elementList = elementService.selectListByForm(form, BaseElement.class);

		ListedHashTree rootTree = new ListedHashTree();
		if(!CollectionUtils.isEmpty(elementList)){
			List<BaseElement> threadGroupList = this.findElements(elementList, testPlanElement.getId());
			if(!CollectionUtils.isEmpty(threadGroupList)){
				this.recursionHashtree(rootTree, testPlanElement, elementList);
				this.applicationContext.publishEvent(new BuildTreeEvent(testPlanElement.getTestId()));
			}
		}
		this.mergeHeader(rootTree);
		//保存测试配置文件
		try {
			JMeterUtils.setProperty("saveservice_properties", JMETER_PATH + "/saveservice.properties");
			SaveService.saveTree(rootTree, new FileOutputStream(JMETER_PATH + "/jmx/" + testPlanElement.getPlanId() + ".jmx"));
		} catch (Exception e) {
			e.printStackTrace();
		}
		return rootTree;
	}

	private void recursionHashtree(ListedHashTree parentTree, BaseElement parent, List<BaseElement> elements) throws Exception{
		TestElement parentTest = this.buildChildElement(parent);
		if(parentTest.isEnabled()){
			List<BaseElement> children;
			if(parentTest instanceof ModuleControllerExt){
				ModuleControllerExt ext = (ModuleControllerExt) parentTest;
				children = ext.getChildren();
			}else{
				children = this.findElements(elements, parent.getId());
			}
			if(children.size()>0){
				ListedHashTree childTree = new ListedHashTree();
				for(BaseElement el : children){
					el.setUserId(parent.getUserId());
					el.setTestId(parent.getTestId());
					this.recursionHashtree(childTree, el, elements);
				}
				parentTree.add(parentTest, childTree);
			}else{
				parentTree.add(parentTest);
			}
		}
	}

	/**
	 * 过滤子元素
	 * @param elements 所有元素
	 * @param parentId 父ID
	 * @return
	 */
	private List<BaseElement> findElements(List<BaseElement> elements, String parentId){
		return elements.stream().filter(e -> e.getParentId().equals(parentId))
				.sorted(Comparator.comparingInt(BaseElement::getSort))
				.collect(Collectors.toList());
	}

	/**
	 * 创建TestElement元件
	 * @param el
	 * @return
	 * @throws Exception
	 */
	private TestElement buildChildElement(BaseElement el) throws Exception{
		ElementBuildService buildService = this.applicationContext.getBean(el.getCategory() + "Impl", ElementBuildService.class);
		return buildService.build(el);
	}

	/**
	 * 合并HttpHeader
	 * @param rootTree
	 */
	private void mergeHeader(ListedHashTree rootTree){
		SearchByClass<HttpHeaderConfig> searcherHeader = new SearchByClass<>(HttpHeaderConfig.class);
		rootTree.traverse(searcherHeader);
		CollectionProperty headerProperties = new CollectionProperty();
		Iterator<HttpHeaderConfig> iterHeader = searcherHeader.getSearchResults().iterator();
		while (iterHeader.hasNext()){
			HeaderManager lopManager = iterHeader.next().getHeaderManager();
			if(lopManager!=null){
				CollectionProperty properties = lopManager.getHeaders();
				PropertyIterator iter = properties.iterator();
				while (iter.hasNext()){
					headerProperties.addProperty(iter.next());
				}
			}
		}
		SearchByClass<HttpRequestSampler> searcherHttp = new SearchByClass<>(HttpRequestSampler.class);
		rootTree.traverse(searcherHttp);
		Iterator<HttpRequestSampler> iterHttp = searcherHttp.getSearchResults().iterator();
		while (iterHttp.hasNext()){
			HttpRequestSampler http = iterHttp.next();
			HeaderManager httpHeader = http.getHeaderManager();
			if(httpHeader==null){
				httpHeader = new HeaderManager();
			}
			PropertyIterator iter = headerProperties.iterator();
			while (iter.hasNext()){
				JMeterProperty p = iter.next();
				httpHeader.add((Header) p.getObjectValue());
			}
			http.setHeaderManager(httpHeader);
		}
	}
}
